#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
使用 3‑点 Lagrange（二次）实现照片放大四倍（128×128 → 512×512）
"""

from PIL import Image
import numpy as np

# -------------------------------------------------
# 1️⃣  通用 Lagrange 插值（单维）函数
# -------------------------------------------------
def lagrange_interp_1d(pixels, scale):
    """
    对一维像素序列进行二次 Lagrange 插值（每段使用相邻 3 点）。
    pixels : 长度为 N 的 1‑D ndarray，元素可以是标量（灰度）或向量（RGB）
    scale  : 放大倍数（这里固定为 4）
    返回   : 长度为 (N‑1)*scale+1 的 ndarray，已完成插值
    """
    N = len(pixels)

    # 为了在首尾也能使用 3 点，复制首尾各一个点
    extended = np.concatenate(([pixels[0]], pixels, [pixels[-1]]), axis=0)  # shape (N+2, …)

    out_len = (N - 1) * scale + 1
    if pixels.ndim == 2:                     # RGB
        out = np.empty((out_len, pixels.shape[1]), dtype=np.float64)
    else:                                    # 灰度
        out = np.empty(out_len, dtype=np.float64)

    # 把原始采样点拷贝到对应位置
    out[::scale] = pixels

    # 对每两个相邻原始点之间插值（使用它们各自左/右的邻点形成 3‑点基组）
    for i in range(N - 1):
        # 基点在 extended 中的索引：i、i+1、i+2 对应原始的 i‑1、i、i+1（首尾已补）
        x_vals = np.array([i, i + 1, i + 2], dtype=np.float64)   # 归一化坐标 0,1,2
        y_vals = extended[i:i + 3]                               # 3 个基点

        for k in range(1, scale):
            t = k / scale               # 归一化位置 0<t<1，实际坐标 = i + t
            x = i + t

            # Lagrange 基函数（二次）
            L0 = (x - x_vals[1]) * (x - x_vals[2]) / ((x_vals[0] - x_vals[1]) * (x_vals[0] - x_vals[2]))
            L1 = (x - x_vals[0]) * (x - x_vals[2]) / ((x_vals[1] - x_vals[0]) * (x_vals[1] - x_vals[2]))
            L2 = (x - x_vals[0]) * (x - x_vals[1]) / ((x_vals[2] - x_vals[0]) * (x_vals[2] - x_vals[1]))

            out[i * scale + k] = L0 * y_vals[0] + L1 * y_vals[1] + L2 * y_vals[2]

    return out

# -------------------------------------------------
# 2️⃣  主放大函数（先横向后纵向）
# -------------------------------------------------
def upscale_lagrange_4x(img):
    """
    将输入图像放大四倍，使用 3‑点 Lagrange（二次）插值。
    参数
    ----
    img : PIL.Image（已转为 RGB 或 L）
    返回
    ----
    out_img : 放大后的 PIL.Image（尺寸 = 4·原尺寸）
    """
    src_w, src_h = img.size
    src_arr = np.asarray(img, dtype=np.float64)   # (h, w, c) 或 (h, w)

    # ---------- 水平插值 ----------
    if src_arr.ndim == 3:          # RGB
        rows_interp = [lagrange_interp_1d(row, 4) for row in src_arr]   # (w*4-3, 3)
        horiz = np.stack(rows_interp, axis=0)                         # (h, w*4-3, 3)
    else:                           # 灰度
        rows_interp = [lagrange_interp_1d(row, 4) for row in src_arr]
        horiz = np.stack(rows_interp, axis=0)                         # (h, w*4-3)

    # ---------- 纵向插值 ----------
    if horiz.ndim == 3:
        trans = np.transpose(horiz, (1, 0, 2))          # (w', h, c)
        cols_interp = [lagrange_interp_1d(col, 4) for col in trans]   # (h*4-3, 3)
        vert = np.stack(cols_interp, axis=0)                         # (w'*4-3, h*4-3, 3)
        final = np.transpose(vert, (1, 0, 2))                         # (h*4-3, w*4-3, 3)
    else:
        trans = np.transpose(horiz)                # (w', h)
        cols_interp = [lagrange_interp_1d(col, 4) for col in trans]
        vert = np.stack(cols_interp, axis=0)                         # (w'*4-3, h*4-3)
        final = np.transpose(vert)                                   # (h*4-3, w*4-3)

    # ---------- 补齐到完整 4 倍尺寸 ----------
    dst_h, dst_w = src_h * 4, src_w * 4
    pad_h = dst_h - final.shape[0]
    pad_w = dst_w - final.shape[1]

    if pad_h > 0:
        final = np.vstack([final, np.repeat(final[-1:, :], pad_h, axis=0)])
    if pad_w > 0:
        final = np.hstack([final, np.repeat(final[:, -1:], pad_w, axis=1)])

    # ---------- 转回 Pillow 图像 ----------
    final = np.clip(final, 0, 255).astype(np.uint8)
    out_img = Image.fromarray(final, mode=img.mode)
    return out_img

# -------------------------------------------------
# 3️⃣  示例入口
# -------------------------------------------------
if __name__ == "__main__":
    input_path  = "lena_128.png"          # 128×128 原图
    output_path = "lena_512_lagrange.png"

    img = Image.open(input_path).convert("RGB")
    upscaled = upscale_lagrange_4x(img)
    upscaled.save(output_path)
    print(f"已完成 4 倍二次 Lagrange 插值放大，保存为 {output_path}")
